//
// Created by jyc on 2021/10/2.
//

#ifndef BFALCOR_DIRECTEDGRAPH_H
#define BFALCOR_DIRECTEDGRAPH_H
#include "PreStd.h"

#include "Logger.h"
namespace Bonnie{
    class DirectedGraph{
    public:
        static constexpr uint32_t kInvalidID = (uint32_t)-1;
        class Node{
            friend class DirectedGraph;
            std::vector<uint32_t> mIncomeEdges;
            std::vector<uint32_t> mOutcomeEdges;
        public:

            Node() = default;
            uint32_t getIncomeEdgeCount()const{return mIncomeEdges.size();}
            uint32_t getOutcomeEdgeCount()const{return mOutcomeEdges.size();}

            uint32_t getIncomeEdge(uint32_t i)const{return mIncomeEdges[i];}
            uint32_t getOutcomeEdge(uint32_t i)const{return mOutcomeEdges[i];}


        };
        class Edge{
            friend class DirectedGraph;
            uint32_t mSrc = kInvalidID;
            uint32_t mDst = kInvalidID;
        public:
            Edge() = default;
            Edge(uint32_t mSrc, uint32_t mDst) : mSrc(mSrc), mDst(mDst) {}

            uint32_t getSourceNode()const {return  mSrc;}
            uint32_t getDestNode()const{return mDst;}
        };

        uint32_t addNode(){
            mNodes[mCurrentNodeIdx] = Node();
            return mCurrentNodeIdx ++;
        }

        /**
         * remove a node from the graph,it also remove all edges associated with the node
         * @param i node id
         * @return a list of removed edges
         */
        std::unordered_set<uint32_t> removeNode(uint32_t i){
            Bonnie_Assert(mNodes.find(i) != mNodes.end(),"the one need to remove doesn`t exist");

            auto& node = mNodes[i];
            std::unordered_set<uint32_t> removedEdges;
            for(auto in:node.mIncomeEdges){
                removeEdgeFromNode<false>(in,mEdges[in].mSrc);
                removedEdges.insert(in);
            }
            for(auto out:node.mOutcomeEdges){
                removeEdgeFromNode<true>(out,mEdges[out].mDst);
                removedEdges.insert(out);
            }
            for(auto eId:removedEdges) mEdges.erase(eId);
            mNodes.erase(i);
            return removedEdges;
        }

        uint32_t addEdge(uint32_t src,uint32_t dst){
            Bonnie_Assert(mNodes.find(src) != mNodes.end() && mNodes.find(dst) != mNodes.end(),"src or dst doesn`t exist");


            mNodes[src].mOutcomeEdges.push_back(mCurrentEdgeIdx);
            mNodes[dst].mIncomeEdges.push_back(mCurrentEdgeIdx);
            mEdges[mCurrentEdgeIdx] = {src, dst};
            return mCurrentEdgeIdx ++;
        }
        void removeEdge(uint32_t i){
            Bonnie_Assert(mEdges.find(i)!= mEdges.end(),"the edge doesn`t exist");

            auto& edge = mEdges[i];
            removeEdgeFromNode<true>(i,edge.mDst);
            removeEdgeFromNode<false>(i,edge.mSrc);
        }

        /** Check if a node exists
*/
        bool doesNodeExist(uint32_t nodeId) const { return mNodes.find(nodeId) != mNodes.end(); }

        /** Check if an edge exists
        */
        bool doesEdgeExist(uint32_t edgeId) const { return mEdges.find(edgeId) != mEdges.end(); }
        /** Get a node
        */
        const Node* getNode(uint32_t nodeId) const
        {
            if (doesNodeExist(nodeId) == false)
            {
                LogWarning("DirectGraph::getNode() - node ID doesn't exist");
                return nullptr;
            }
            return &mNodes.at(nodeId);
        }

        /** Get an edge
        */
        const Edge* getEdge(uint32_t edgeId)
        {
            if (doesEdgeExist(edgeId) == false)
            {
                LogWarning("DirectGraph::getEdge() - edge ID doesn't exist");
                return nullptr;
            }
            return &mEdges[edgeId];
        }
        uint32_t getCurrentNodeId() const { return mCurrentNodeIdx; }
        uint32_t getCurrentEdgeId() const { return mCurrentEdgeIdx; }
    private:

        uint32_t mCurrentNodeIdx = 0;
        uint32_t mCurrentEdgeIdx = 0;
        std::unordered_map<uint32_t ,Node>  mNodes;
        std::unordered_map<uint32_t, Edge> mEdges;

        template<bool removeInput>
        void removeEdgeFromNode(uint32_t eId,uint32_t nId){
            auto& vec = removeInput?mNodes[nId].mIncomeEdges:mNodes[nId].mOutcomeEdges;
            for(auto e = vec.begin();e != vec.end();e++){
                if(*e == eId){
                    vec.erase(e);
                    return;
                }
            }
            Bonnie_Should_Not_Be_Here();
        }

    };
}
#endif //BFALCOR_DIRECTEDGRAPH_H
